{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "911b3b37-3b29-4833-94f2-bfe47af00c83",
   "metadata": {},
   "source": [
    "# Lesson 6: Essay Writer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5762271-8736-4e94-9444-8c92bd0e8074",
   "metadata": {
    "height": 63
   },
   "outputs": [],
   "source": [
    "from dotenv import load_dotenv\n",
    "\n",
    "_ = load_dotenv()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0168aee-bce9-4d60-b827-f86a88187e31",
   "metadata": {
    "height": 148
   },
   "outputs": [],
   "source": [
    "from langgraph.graph import StateGraph, END\n",
    "from typing import TypedDict, Annotated, List\n",
    "import operator\n",
    "from langgraph.checkpoint.sqlite import SqliteSaver\n",
    "from langchain_core.messages import AnyMessage, SystemMessage, HumanMessage, AIMessage, ChatMessage\n",
    "\n",
    "memory = SqliteSaver.from_conn_string(\":memory:\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2589c5b6-6cc2-4594-9a17-dccdcf676054",
   "metadata": {
    "height": 148
   },
   "outputs": [],
   "source": [
    "class AgentState(TypedDict):\n",
    "    task: str\n",
    "    plan: str\n",
    "    draft: str\n",
    "    critique: str\n",
    "    content: List[str]\n",
    "    revision_number: int\n",
    "    max_revisions: int"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2ba84ec-c172-4de7-ac55-e3158a531b23",
   "metadata": {
    "height": 46
   },
   "outputs": [],
   "source": [
    "from langchain_openai import ChatOpenAI\n",
    "model = ChatOpenAI(model=\"gpt-3.5-turbo\", temperature=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "876d5092-b8ef-4e38-b4d7-0e80c609bf7a",
   "metadata": {
    "height": 80
   },
   "outputs": [],
   "source": [
    "PLAN_PROMPT = \"\"\"You are an expert writer tasked with writing a high level outline of an essay. \\\n",
    "Write such an outline for the user provided topic. Give an outline of the essay along with any relevant notes \\\n",
    "or instructions for the sections.\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10084a02-2928-4945-9f7c-ad3f5b33caf7",
   "metadata": {
    "height": 165
   },
   "outputs": [],
   "source": [
    "WRITER_PROMPT = \"\"\"You are an essay assistant tasked with writing excellent 5-paragraph essays.\\\n",
    "Generate the best essay possible for the user's request and the initial outline. \\\n",
    "If the user provides critique, respond with a revised version of your previous attempts. \\\n",
    "Utilize all the information below as needed: \n",
    "\n",
    "------\n",
    "\n",
    "{content}\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "714d1205-f8fc-4912-b148-2a45da99219c",
   "metadata": {
    "height": 80
   },
   "outputs": [],
   "source": [
    "REFLECTION_PROMPT = \"\"\"You are a teacher grading an essay submission. \\\n",
    "Generate critique and recommendations for the user's submission. \\\n",
    "Provide detailed recommendations, including requests for length, depth, style, etc.\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83588e70-254f-4f83-a510-c8ae81e729b0",
   "metadata": {
    "height": 97
   },
   "outputs": [],
   "source": [
    "RESEARCH_PLAN_PROMPT = \"\"\"You are a researcher charged with providing information that can \\\n",
    "be used when writing the following essay. Generate a list of search queries that will gather \\\n",
    "any relevant information. Only generate 3 queries max.\"\"\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6cb3ef4c-58b3-401b-b104-0d51e553d982",
   "metadata": {
    "height": 97
   },
   "outputs": [],
   "source": [
    "RESEARCH_CRITIQUE_PROMPT = \"\"\"You are a researcher charged with providing information that can \\\n",
    "be used when making any requested revisions (as outlined below). \\\n",
    "Generate a list of search queries that will gather any relevant information. Only generate 3 queries max.\"\"\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc3293b7-a50c-43c8-a022-8975e1e444b8",
   "metadata": {
    "height": 80
   },
   "outputs": [],
   "source": [
    "from langchain_core.pydantic_v1 import BaseModel\n",
    "\n",
    "class Queries(BaseModel):\n",
    "    queries: List[str]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0722c3d4-4cbf-43bf-81b0-50f634c4ce61",
   "metadata": {
    "height": 63
   },
   "outputs": [],
   "source": [
    "from tavily import TavilyClient\n",
    "import os\n",
    "tavily = TavilyClient(api_key=os.environ[\"TAVILY_API_KEY\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b2f82fe-3ec4-4917-be51-9fb10d1317fa",
   "metadata": {
    "height": 131
   },
   "outputs": [],
   "source": [
    "def plan_node(state: AgentState):\n",
    "    messages = [\n",
    "        SystemMessage(content=PLAN_PROMPT), \n",
    "        HumanMessage(content=state['task'])\n",
    "    ]\n",
    "    response = model.invoke(messages)\n",
    "    return {\"plan\": response.content}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee0fe1c7-77e2-499c-a2f9-1f739bb6ddf0",
   "metadata": {
    "height": 199
   },
   "outputs": [],
   "source": [
    "def research_plan_node(state: AgentState):\n",
    "    queries = model.with_structured_output(Queries).invoke([\n",
    "        SystemMessage(content=RESEARCH_PLAN_PROMPT),\n",
    "        HumanMessage(content=state['task'])\n",
    "    ])\n",
    "    content = state['content'] or []\n",
    "    for q in queries.queries:\n",
    "        response = tavily.search(query=q, max_results=2)\n",
    "        for r in response['results']:\n",
    "            content.append(r['content'])\n",
    "    return {\"content\": content}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98f303b1-a4d0-408c-8cc0-515ff980717f",
   "metadata": {
    "height": 301
   },
   "outputs": [],
   "source": [
    "def generation_node(state: AgentState):\n",
    "    content = \"\\n\\n\".join(state['content'] or [])\n",
    "    user_message = HumanMessage(\n",
    "        content=f\"{state['task']}\\n\\nHere is my plan:\\n\\n{state['plan']}\")\n",
    "    messages = [\n",
    "        SystemMessage(\n",
    "            content=WRITER_PROMPT.format(content=content)\n",
    "        ),\n",
    "        user_message\n",
    "        ]\n",
    "    response = model.invoke(messages)\n",
    "    return {\n",
    "        \"draft\": response.content, \n",
    "        \"revision_number\": state.get(\"revision_number\", 1) + 1\n",
    "    }\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf4dcb93-6298-4cfd-b3ce-61dfac7fb35f",
   "metadata": {
    "height": 131
   },
   "outputs": [],
   "source": [
    "def reflection_node(state: AgentState):\n",
    "    messages = [\n",
    "        SystemMessage(content=REFLECTION_PROMPT), \n",
    "        HumanMessage(content=state['draft'])\n",
    "    ]\n",
    "    response = model.invoke(messages)\n",
    "    return {\"critique\": response.content}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "932883a4-c722-42bb-aec0-b4f41c5c81a4",
   "metadata": {
    "height": 199
   },
   "outputs": [],
   "source": [
    "def research_critique_node(state: AgentState):\n",
    "    queries = model.with_structured_output(Queries).invoke([\n",
    "        SystemMessage(content=RESEARCH_CRITIQUE_PROMPT),\n",
    "        HumanMessage(content=state['critique'])\n",
    "    ])\n",
    "    content = state['content'] or []\n",
    "    for q in queries.queries:\n",
    "        response = tavily.search(query=q, max_results=2)\n",
    "        for r in response['results']:\n",
    "            content.append(r['content'])\n",
    "    return {\"content\": content}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff362f49-dcf1-4ea1-a86c-e516e9ab897d",
   "metadata": {
    "height": 80
   },
   "outputs": [],
   "source": [
    "def should_continue(state):\n",
    "    if state[\"revision_number\"] > state[\"max_revisions\"]:\n",
    "        return END\n",
    "    return \"reflect\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7e15a20-83d7-434c-8551-bce8dcc32be0",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "builder = StateGraph(AgentState)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54ab2c74-f32e-490c-a85d-932d11444210",
   "metadata": {
    "height": 97
   },
   "outputs": [],
   "source": [
    "builder.add_node(\"planner\", plan_node)\n",
    "builder.add_node(\"generate\", generation_node)\n",
    "builder.add_node(\"reflect\", reflection_node)\n",
    "builder.add_node(\"research_plan\", research_plan_node)\n",
    "builder.add_node(\"research_critique\", research_critique_node)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a833d3ce-bd31-4319-811d-decff226b970",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "builder.set_entry_point(\"planner\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76e93cce-6eab-4c7c-ac64-e9993fdb30d6",
   "metadata": {
    "height": 114
   },
   "outputs": [],
   "source": [
    "builder.add_conditional_edges(\n",
    "    \"generate\", \n",
    "    should_continue, \n",
    "    {END: END, \"reflect\": \"reflect\"}\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd2d0990-a932-423f-9ff3-5cada58c5f32",
   "metadata": {
    "height": 97
   },
   "outputs": [],
   "source": [
    "builder.add_edge(\"planner\", \"research_plan\")\n",
    "builder.add_edge(\"research_plan\", \"generate\")\n",
    "\n",
    "builder.add_edge(\"reflect\", \"research_critique\")\n",
    "builder.add_edge(\"research_critique\", \"generate\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27cde654-64e2-48bc-80a9-0ed668ccb7dc",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "graph = builder.compile(checkpointer=memory)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4871f644-b131-4065-b7ce-b82c20a41f11",
   "metadata": {
    "height": 63
   },
   "outputs": [],
   "source": [
    "from IPython.display import Image\n",
    "\n",
    "Image(graph.get_graph().draw_png())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98f3be1d-cc4c-41fa-9863-3e386e88e305",
   "metadata": {
    "height": 148
   },
   "outputs": [],
   "source": [
    "thread = {\"configurable\": {\"thread_id\": \"1\"}}\n",
    "for s in graph.stream({\n",
    "    'task': \"what is the difference between langchain and langsmith\",\n",
    "    \"max_revisions\": 2,\n",
    "    \"revision_number\": 1,\n",
    "}, thread):\n",
    "    print(s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ad8a6cc-65d4-4ce7-87aa-4e67d7c23d7b",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "4d1664b5-75e0-46b7-9c2b-4ac9171f4597",
   "metadata": {},
   "source": [
    "## Essay Writer Interface"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e0ae270-3ec3-484a-b729-df7d2b7b0f76",
   "metadata": {
    "height": 80
   },
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "from helper import ewriter, writer_gui"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0ebfa79-c7fc-4aaa-b668-64e5b6cede80",
   "metadata": {
    "height": 63
   },
   "outputs": [],
   "source": [
    "MultiAgent = ewriter()\n",
    "app = writer_gui(MultiAgent.graph)\n",
    "app.launch()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "592b5e62-a203-433c-92a0-3783f490cde1",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14fa923c-7e4f-42d1-965f-0f8ccd50fbd7",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "570c6245-2837-4ac5-983b-95f61f3ac10d",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b910915-b087-4d35-afff-0ec30a5852f1",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4feb6cc-5129-4a99-bb45-851bc07b5709",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e85a02b4-96cc-4b01-8792-397a774eb499",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae8b86a6-5e20-4252-b1d8-009b8318345a",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af925917-b746-48c9-ac74-62fefbe5246c",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5048f2c-4d82-49a5-9cb1-918d78b39f7b",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "393f7f1f-68b4-4462-bfa5-b6472ef1304a",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43ac0aa9-baa7-4b58-889d-2118cc00c6b5",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed6098b9-e2a9-4767-8cb5-346db835c8d2",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d23cf2a-a179-44dc-9ae3-2eddda4b67b4",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14a6005b-0221-4f5e-9be0-0580c1d03126",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41c1ec12-f1c8-41ae-bb3e-5f28997b9b99",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c8c07d7-be17-4c17-82c5-6fe1db028b8b",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04592c8e-1cfe-4b26-93b5-caf1ed1e7d24",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6181c4a9-0e71-4f67-b71f-18a225e37202",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1c478a9-7bfe-49e2-8a7d-1536271f45a6",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a6d6771-3fad-4f37-9b32-45b36ad85c59",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3629eb3-655d-467a-b413-63f547c2de08",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f772f251-2b61-4d10-97c5-61cef9207a76",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0de92979-7ac5-4a7c-91c1-10806b7d529c",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "479c4325-f625-4bbf-9d74-cc58f10763f2",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4070be7-72da-42f9-a25d-8a6c628788b8",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9289efbe-7033-4f32-8482-2039c5f9db90",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25e480bb-22ab-4acb-a42c-71da3d04a5b1",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90dea35c-7483-4b3d-b5e3-76eb3a0fe536",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9ac5730-a9d5-4ea4-8546-ebcb265cf1da",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96e1f28b-46d8-4bcd-b2e4-730376ee7ccf",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22ac7020-b4f4-4bd2-a875-ccee93f83d83",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61f79eb9-d1c9-44b0-9efd-a8f9b380332a",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce509206-bde1-43e4-a88f-8a565539d357",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdba1590-9e7b-4c0f-9492-81a07d286c55",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa8fe4a8-5372-479d-b248-af7a295c86c1",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7514720a-14bc-4552-ade5-fa03f86f4c73",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
